#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/VectorField.h"
#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/DTI.h"

#include <string>

#include <omp.h>

using namespace std;

char *pDTIPathname = NULL;
char *pWMMaskPathname = NULL;
char *pVectorFieldPathname = NULL;
char *pFiberPathname = NULL;


float stepSize = 0.1f;
float maxLength = 50.0f;
float faThreshold = 0.001f;
float angleThreshold = 0.866f;
int scale = 1;

CVectorField *pDTIVec = NULL;
CVectorField *pVectorField = NULL;
CScalar *pWMMask = NULL;

void PrintUsage()
{
	printf("Usage: ");
	printf("FiberTracking -vec <VectorField> -dti <DTI> -wm_mask <WMMask> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
	printf("-vec <VectorField> = input VectorField nrrd file\n");
	printf("-wm_mask <WMMask> = input White Matter mask nrrd file\n");
	printf("-fiber <Fiber> = output Fiber configure file\n");
	printf("-stepsize <stepsize> = step size\n");
	printf("-maxlength <maxlength> = maximum fiber length\n");
	printf("-fa_threshold <fa_thresold> = FA threshold used to differ white matter and grey matter\n");
	printf("-angle_threshold <angle_thresold> = angle threshold between two directions, in cosine\n");
	printf("-scale <scale> = scale of the result\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-vec") == 0) {
			pVectorFieldPathname = argv[++i];
		} else if (strcmp(argv[i], "-wm_mask") == 0) {
			pWMMaskPathname = argv[++i];
		} else if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) {
			scale = (uint)atoi(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pVectorFieldPathname == NULL || pWMMaskPathname == NULL || pFiberPathname == NULL)
		return -2;

	return 0;
}


int GetNextDirection(CPoint3F *pos, CVector3F *in, CVector3F *out)
{
	if ((pWMMask->GetScalar(pos) > 0.4f || pWMMask->GetScalar(pos) < 0.2f))
		return 0;

	CVector3F vecs[3];
	int count = pVectorField->GetVectors(pos, vecs);
	if (count == 0) {
		return 0;
	}
	
	/* find best output direction */
	float max_v = 0;
	for (int i = 0; i < count; ++i) {
		float v = InnerProduct(in, &(vecs[i]));
		if (fabs(v) > fabs(max_v)) {
			max_v = v;
			out->m_x = vecs[i].m_x;
			out->m_y = vecs[i].m_y;
			out->m_z = vecs[i].m_z;
		}
	}

	if (fabs(max_v) < angleThreshold) {
		return 0;
	}

	if (max_v < 0.0f) {
		out->m_x = -out->m_x;
		out->m_y = -out->m_y;
		out->m_z = -out->m_z;
	}

	return 1;
}

int TrackingFibers(CFiber **fibers, CPoint3F pos)
{
	CVector3F dirs[4];
	int count = 0;
	if ((pWMMask->GetScalar(pos) > 0.4f || pWMMask->GetScalar(pos) < 0.2f)) {
		dirs[0] = CVector3F(0, 0, 0);
		count = 1;
	} else {
		count = pVectorField->GetVectors(pos, dirs);
		count = 1;
	}
	
	*fibers = new CFiber[count];

	for (int i = 0; i < count; ++i) {
		Seed seed;
		seed.pos = pos;
		seed.dir = dirs[i];

		int innerSteps = 10;
		uint maxSteps = (uint)(maxLength / (stepSize * (float)innerSteps));
		uint fnum, bnum;
		CPoint3F *pF = new CPoint3F[maxSteps];
		CPoint3F *pB = new CPoint3F[maxSteps];
		CPoint3F lastPos;
		CVector3F lastDir;

		/* tracking forward */
		lastPos = seed.pos;
		lastDir = seed.dir;
		pF[0] = lastPos;
		for (fnum = 1; fnum < maxSteps; ++fnum) {
			CPoint3F nextPos;
			CVector3F nextDir;

			bool flag = true;
			for (int step = 0; step < innerSteps; ++step) {
				nextPos = lastPos + lastDir * stepSize;

				if (GetNextDirection(&nextPos, &lastDir, &nextDir) == 0) {
					flag = false;
					break;
				}
				lastDir = nextDir;
				lastPos = nextPos;
			}
			if (flag == false)
				break;

			pF[fnum] = nextPos;
			lastDir = nextDir;
			lastPos = nextPos;
		}

		/* tracking backward */
		lastPos = seed.pos;
		lastDir = seed.dir;
		pB[0] = lastPos;
		for (bnum = 1; bnum < maxSteps; ++bnum) {
			CPoint3F nextPos;
			CVector3F nextDir;

			bool flag = true;
			for (int step = 0; step < innerSteps; ++step) {
				nextPos = lastPos - lastDir * stepSize;

				if (GetNextDirection(&nextPos, &lastDir, &nextDir) == 0) {
					flag = false;
					break;
				}
				lastDir = nextDir;
				lastPos = nextPos;
			}
			if (flag == false)
				break;

			pB[bnum] = nextPos;
			lastDir = nextDir;
			lastPos = nextPos;
		}

		(*fibers)[i].CreateFiber(pF, pB, fnum, bnum, &seed);

		SafeDeleteArray(pF);
		SafeDeleteArray(pB);
	}

	return count;
}


void FiberTracking()
{
	uint w = pVectorField->GetWidth();
	uint h = pVectorField->GetHeight();
	uint d = pVectorField->GetDepth();

	int sw = (int)w * scale;
	int sh = (int)h * scale;
	int sd = (int)d * scale;

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, stepSize, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	omp_set_num_threads(64);

	for (int y = 50; y < 51; ++y) {
		printf("Fiber tracking on slice y = %d ... ", y);

		CFiber **fibers = new CFiber*[sw*sd];
		int *count = new int[sw*sd];
		memset(fibers, 0, sizeof(CFiber*)*sw*sd);
		memset(count, 0, sizeof(int)*sw*sd);

		for (int z = 0; z < sd; ++z) {
#pragma omp parallel for
			for (int x = 0; x < sw; ++x) {
				CPoint3F pos = CPoint3F();
				pos.m_x = (float)x/(float)scale + 1.0f / ((float)scale * 2.0f);
				pos.m_y = (float)y/(float)scale + 1.0f / ((float)scale * 2.0f);
				pos.m_z = (float)z/(float)scale + 1.0f / ((float)scale * 2.0f);
				int index = z * sw + x;
				int temp_index = (z * sh + y) * sw + x;
				count[index] = TrackingFibers(&(fibers[index]), pos);
			}
			float ff = float(z+1) / float(sd) * 100.0f;
			printf("\rFiber tracking on slice y = %d ... %.3f%%", y, ff);
		}

		/* Save fibers */
		SaveFibersAsVTK("C:\\Ziang\\Dataset\\Alex\\2X\\ForFiberTracking\\test.vtk", fibers, count, sw, sd);
      
		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_y=%08.04f.conf", path.c_str(), (float)y/(float)scale);
		sprintf(vertPathname, "%sslice_y=%08.04f.vert", path.c_str(), (float)y/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

		//SaveFibers(confPathname, vertPathname, fibers, count, sw, sd);

		for (uint i = 0; i < sd*sw; ++i) {
			SafeDeleteArray(fibers[i]);
		}
		SafeDeleteArray(fibers);
		SafeDeleteArray(count);

        printf("Done.\n");
	}

	fclose(fp);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTIVec = new CVectorField();
	pDTIVec->OpenVectorFieldFile(pDTIPathname);
	pVectorField = new CVectorField();
	pVectorField->OpenVectorFieldFile(pVectorFieldPathname);
	pWMMask = new CScalar();
	pWMMask->OpenScalarFile(pWMMaskPathname);

	FiberTracking();
	
	SafeDelete(pDTIVec);
	SafeDelete(pVectorField);
	SafeDelete(pWMMask);

	return 0;
}
